import io
import threading
from flask import Flask, render_template, Response, request, redirect, url_for
import os
import uuid
import cv2
import torch
from werkzeug.utils import secure_filename
from deep_sort_realtime.deepsort_tracker import DeepSort
from py.login import check_user
from py.users import create_user_table, get_all_users, create_user_table_html

app = Flask(__name__)

UPLOAD_FOLDER = 'uploads'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# 1) 加载 YOLOv5 模型（这里使用 ultralytics 库）
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
model.classes = [0]  # 只保留“人”类别 (COCO 中 class 0 为 person)

# 2) 初始化 DeepSORT 跟踪器
tracker = DeepSort(max_age=30)

# 3) 打开摄像头
cap = cv2.VideoCapture(0)
lock = threading.Lock()

def gen_frames():
    while True:
        success, frame = cap.read()
        if not success:
            break

        results = model(frame)
        detections = results.xyxy[0].cpu().numpy()

        dets_for_tracker = []
        for *xyxy, conf, cls in detections:
            x1, y1, x2, y2 = map(int, xyxy)
            w, h = x2 - x1, y2 - y1
            bbox = [x1, y1, w, h]
            dets_for_tracker.append((bbox, float(conf), int(cls)))

        tracks = tracker.update_tracks(dets_for_tracker, frame=frame)

        for track in tracks:
            if not track.is_confirmed():
                continue
            x1, y1, x2, y2 = map(int, track.to_ltrb())
            track_id = track.track_id
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0,255,0), 2)
            cv2.putText(frame, f'ID {track_id}', (x1, y1-10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)

        ret, buffer = cv2.imencode('.jpg', frame)
        frame_bytes = buffer.tobytes()
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')

@app.route('/mb')
def index():
    return render_template('mb.html')

@app.route('/video_feed')
def video_feed():
    return Response(gen_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')

# ================= FGX 登录和用户管理 =================

@app.route('/')
def hello_world():
    return render_template('login.html')

@app.route('/login', methods=["POST"])
def login():
    username = request.form.get('username')
    password = request.form.get('password')
    key = check_user(username, password)
    if key:
        return render_template('index.html')
    else:
        return render_template('login_error.html')

@app.route('/users')
def user():
    create_user_table()
    users = get_all_users()
    user_table = create_user_table_html(users)
    return render_template('user.html', user_table=user_table)

# ================= 视频文件上传追踪 =================

@app.route('/mbb')
def upload_page():
    return render_template('mb2.html')

@app.route('/upload', methods=['POST'])
def upload():
    if 'file' not in request.files:
        return 'No file part'
    file = request.files['file']
    if file.filename == '':
        return 'No selected file'

    filename = secure_filename(file.filename)
    filepath = os.path.join(app.config['UPLOAD_FOLDER'], f"{uuid.uuid4()}_{filename}")
    file.save(filepath)
    return redirect(url_for('track_video', filename=os.path.basename(filepath)))

def generate_tracking_frames(path):
    cap = cv2.VideoCapture(path)
    while cap.isOpened():
        success, frame = cap.read()
        if not success:
            break

        results = model(frame)
        detections = results.xyxy[0].cpu().numpy()
        dets_for_tracker = []
        for *xyxy, conf, cls in detections:
            x1, y1, x2, y2 = map(int, xyxy)
            w, h = x2 - x1, y2 - y1
            bbox = [x1, y1, w, h]
            dets_for_tracker.append((bbox, float(conf), int(cls)))

        tracks = tracker.update_tracks(dets_for_tracker, frame=frame)

        for track in tracks:
            if not track.is_confirmed():
                continue
            x1, y1, x2, y2 = map(int, track.to_ltrb())
            track_id = track.track_id
            cv2.rectangle(frame, (x1, y1), (x2, y2), (255,0,0), 2)
            cv2.putText(frame, f'ID {track_id}', (x1, y1-10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,0,0), 2)

        ret, buffer = cv2.imencode('.jpg', frame)
        frame_bytes = buffer.tobytes()
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')

@app.route('/track/<filename>')
def track_video(filename):
    filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    return Response(generate_tracking_frames(filepath), mimetype='multipart/x-mixed-replace; boundary=frame')

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, threaded=True)
